import os, json, requests, base64, string, sys
from pathlib import Path
from modelscope import AutoModelForCausalLM, AutoTokenizer
from zlib import compress
from tqdm import tqdm

model_name = "Qwen/Qwen2.5-0.5B-Instruct"
if len(sys.argv) == 2:
    check_point_path = sys.argv[1]
    if check_point_path.strip() != '':
        model_name = f"output/qwen2_5-0_5b-instruct/{check_point_path}"
print('model_name', model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name, torch_dtype="auto", device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)


def gen_plant_uml(prompt):
    messages = [
        {
            "role": "system",
            "content": "你是plantuml代码生成助手。不要输出与代码无关的任何东西。流程使用中文。",
        },
        {"role": "user", "content": prompt},
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=512,
    )
    generated_ids = [
        output_ids[len(input_ids) :]
        for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return response


plantuml_alphabet = (
    string.digits + string.ascii_uppercase + string.ascii_lowercase + "-_"
)
base64_alphabet = string.ascii_uppercase + string.ascii_lowercase + string.digits + "+/"
b64_to_plantuml = bytes.maketrans(
    base64_alphabet.encode("utf-8"), plantuml_alphabet.encode("utf-8")
)


def check_plantuml_code(plantuml_code: str):
    # 太短的直接失败
    if len(plantuml_code.split()) <= 8:
        return False
    
    # 至少有一个分支
    if plantuml_code.count('if') == 0:
        return False
    
    # if endif 数量不一致
    if plantuml_code.count('if') != 2 * (plantuml_code.count('endif') + plantuml_code.count('end if')):
        return False
    
    # 将字符串编码为plantuml格式
    zlibbed_str = compress(plantuml_code.encode("utf-8"))
    compressed_string = zlibbed_str[2:-4]
    plantuml_code_encoded = (
        base64.b64encode(compressed_string).translate(b64_to_plantuml).decode("utf-8")
    )

    resp = requests.get(
        f"http://localhost:13001/plantuml/txt/{plantuml_code_encoded}"
    ).text
    if resp.strip().startswith('[From'):
        return False
    return True

success = 0
fail = 0
res = []
funcs = json.loads(Path("functions.json").read_text())
for func in tqdm(funcs):
    query = f"根据下面的需求生成一段plantuml流程图伪代码,只输出代码无需解释，不要输出与代码无关的任何东西,if和endif要一一对应。这是需求，{func}"
    plantuml_code = gen_plant_uml(query)
    plantuml_code = plantuml_code.replace("```", "")
    plantuml_code = plantuml_code.replace("；", ";")
    plantuml_code = plantuml_code.replace("：", ":")
    plantuml_code = plantuml_code.strip()
    check_result = check_plantuml_code(plantuml_code)
    
    # 显示
    print(f'{check_result} {func}\n:::::\n{plantuml_code}\n---------------------------')
    
    if check_result:
        success += 1
        res.append({"query": query, "response": plantuml_code})
    else:
        fail += 1

# 保存到data.json
data_json: list = json.loads(Path('data.json').read_text())
data_json.extend(res)
Path('data.json').write_text(json.dumps(data_json, ensure_ascii=False))
print(f'success: {success}, fail: {fail}')
